#ifndef HEADER_SMS4_H
#define HEADER_SMS4_H

#define SM4_KEY_LENGTH 16  // 密钥长度
#define SM4_BLOCK_SIZE 16  // 分组长度
#define SM4_IV_LENGTH  (SM4_BLOCK_SIZE)
#define SM4_NUM_ROUNDS 32  // 迭代轮数

#ifdef __KERNEL__
#include <linux/types.h>
#include <linux/string.h>
#else
#include <stdint.h>
#include <string.h>
#include <sys/types.h>
#endif
#ifdef __cplusplus
extern "C" {
#endif

typedef struct {
    uint32_t rk[SM4_NUM_ROUNDS];
} sm4_key_t;

void sm4_set_encrypt_key(sm4_key_t* key, const unsigned char user_key[16]);

void sm4_set_decrypt_key(sm4_key_t* key, const unsigned char user_key[16]);

void sm4_encrypt(const unsigned char in[16], unsigned char out[16],
                 const sm4_key_t* key);

#define sm4_decrypt(in, out, key) sm4_encrypt(in, out, key)

void sm4_ecb_encrypt(const unsigned char* in, unsigned char* out,
                     const sm4_key_t* key, int enc);

void sm4_cbc_encrypt(const unsigned char* in, unsigned char* out, size_t len,
                     const sm4_key_t* key, unsigned char* iv, int enc);

void sm4_cfb128_encrypt(const unsigned char* in, unsigned char* out, size_t len,
                        const sm4_key_t* key, unsigned char* iv, int* num,
                        int enc);

void sm4_ofb128_encrypt(const unsigned char* in, unsigned char* out, size_t len,
                        const sm4_key_t* key, unsigned char* iv, int* num);

void sm4_ctr128_encrypt(const unsigned char* in, unsigned char* out, size_t len,
                        const sm4_key_t* key, unsigned char* iv,
                        unsigned char ecount_buf[SM4_BLOCK_SIZE],
                        unsigned int* num);

int sm4_wrap_key(sm4_key_t* key, const unsigned char* iv, unsigned char* out,
                 const unsigned char* in, unsigned int inlen);

int sm4_unwrap_key(sm4_key_t* key, const unsigned char* iv, unsigned char* out,
                   const unsigned char* in, unsigned int inlen);

void sm4_ctr32_encrypt_blocks(const unsigned char* in, unsigned char* out,
                              size_t blocks, const sm4_key_t* key,
                              const unsigned char iv[16]);

#define SMS4_EDE_KEY_LENGTH (SM4_KEY_LENGTH * 3)

typedef struct {
    sm4_key_t k1;
    sm4_key_t k2;
    sm4_key_t k3;
} sm4_ede_key_t;

void sm4_ede_set_encrypt_key(sm4_ede_key_t*      key,
                             const unsigned char user_key[48]);

void sm4_ede_set_decrypt_key(sm4_ede_key_t*      key,
                             const unsigned char user_key[48]);

void sm4_ede_encrypt(const unsigned char in[16], unsigned char out[16],
                     const sm4_ede_key_t* key);

#define sm4_ede_decrypt(in, out, key) sm4_ede_encrypt(in, out, key)

void sm4_ede_ecb_encrypt(const unsigned char* in, unsigned char* out,
                         const sm4_ede_key_t* key, int enc);

void sm4_ede_cbc_encrypt(const unsigned char* in, unsigned char* out,
                         size_t len, const sm4_ede_key_t* key,
                         unsigned char* iv, int enc);

void sm4_ede_cfb128_encrypt(const unsigned char* in, unsigned char* out,
                            size_t len, const sm4_ede_key_t* key,
                            unsigned char* iv, int* num, int enc);

void sm4_ede_ofb128_encrypt(const unsigned char* in, unsigned char* out,
                            size_t len, const sm4_ede_key_t* key,
                            unsigned char* iv, int* num);

void sm4_ede_ctr128_encrypt(const unsigned char* in, unsigned char* out,
                            size_t len, const sm4_ede_key_t* key,
                            unsigned char* iv,
                            unsigned char  ecount_buf[SM4_BLOCK_SIZE],
                            unsigned int*  num);

int sm4_ede_wrap_key(sm4_ede_key_t* key, const unsigned char* iv,
                     unsigned char* out, const unsigned char* in,
                     unsigned int inlen);

int sm4_ede_unwrap_key(sm4_ede_key_t* key, const unsigned char* iv,
                       unsigned char* out, const unsigned char* in,
                       unsigned int inlen);

#ifdef __cplusplus
}
#endif
#endif
